import glob
import xml.etree.ElementTree as ET
import os


def xml_to_yolo(xml_path, img_path, save_dir):
    img_name = os.path.basename(img_path)
    xml_name_pre = os.path.basename(xml_path).split(".")[0]
    img_name_pre = os.path.basename(img_path).split(".")[0]
    if xml_name_pre != img_name_pre:
        print("xml_name is not equal to img_name")
        return
    tree = ET.parse(xml_path)
    root = tree.getroot()
    # 拼接格式 图片地址 类别1  cx cy w h  类别2  cx cy w h
    img_annotation = img_name
    for obj in root.findall('object'):
        class_name = obj.find('name').text
        xmin = float(obj.find('bndbox/xmin').text)
        ymin = float(obj.find('bndbox/ymin').text)
        xmax = float(obj.find('bndbox/xmax').text)
        ymax = float(obj.find('bndbox/ymax').text)
        # Convert to YOLO format
        x_center = int((xmin + xmax) / 2)
        y_center = int((ymin + ymax) / 2)
        box_width = int((xmax - xmin))
        box_height = int((ymax - ymin))
        cls_id = cls_dic[class_name]
        # cls_id = 0
        img_annotation += f"\t\t{cls_id}\t{x_center}\t{y_center}\t{box_width}\t{box_height}\t"
    # Create a text file to save YOLO annotations
    file_name = os.path.splitext(os.path.basename(xml_path))[0] + '.txt'
    if os.path.isdir(save_directory):
        save_path = os.path.join(save_dir, file_name)
    else:
        save_path = save_dir

    with open(save_path, 'a+') as file:
        file.write(img_annotation + '\n')
    return save_path


if __name__ == '__main__':
    # 保存为YOLOv3需要的txt格式
    save_directory = r'../data/VOC2007/yolo_annotation.txt'
    # 获取转换宽高为416x416之后的标签，用于进行等比例缩放
    xml_paths = glob.glob(os.path.join(r'../data/VOC2007/Annotations', "*"))
    # 转换之后的原图像
    img_paths = glob.glob(os.path.join(r"../data/VOC2007/YOLOv3_JPEGImages", "*"))
    # cls_dic = {"fish_gray": 0, "fish_red": 1, "fish_black": 2}
    # cls_dic = {"person": 0, "dog": 1, "cat": 2, "horse": 3}
    cls_dic = {"big_fish": 0, "small_fish": 1}
    for idx, xml_path in enumerate(xml_paths):
        img_path = img_paths[idx]
        saved_path = xml_to_yolo(xml_path, img_path, save_directory)
        print(f"YOLO annotations saved to: {saved_path}")
